package SecurityMessage;

import RC4.RC4Utils;

import javax.crypto.spec.SecretKeySpec;
import javax.swing.*;
import java.awt.*;
import java.awt.event.ActionEvent;
import java.awt.event.ActionListener;
import java.io.*;
import java.net.ServerSocket;
import java.net.Socket;
import java.nio.charset.StandardCharsets;
import java.security.Key;
import java.security.KeyPair;
import java.security.NoSuchAlgorithmException;
import java.security.PublicKey;
import java.util.HashMap;
import java.util.Map;

/**
 * @创建人 MBY381
 * @创建时间 2022/6/23 16:13
 * @描述 安全通信软件，使用RSA2048进行非对称数据加密（已完成），HashMAC进行消息完整性验证（待添加）
 */
public class ClientThread implements Runnable {
    private ServerSocket serverSocket;
    private JFrame jFrame;
    private static final Font MY_FONT = new Font("宋体", Font.BOLD, 30);
    private static final String SERVER_IP = "127.0.0.1";
    private final String TARGET_IP = "127.0.0.1";
    private static final Integer SERVER_PORT = 8888;
    private static final String RC4_KEY = "mbyzzqwdnmd";

    public KeyPair getMyKeyPair() {
        return myKeyPair;
    }

    public void setMyKeyPair(KeyPair myKeyPair) {
        this.myKeyPair = myKeyPair;
    }

    private Map<Integer, PublicKey> publicKeyMap;
    private KeyPair myKeyPair;
    private Integer myPort;
    private Integer friendPort;
    String[] talkContent = new String[]{" ", " ", " ", " ", " "};
    private Integer contentTag = 0;
    private String myMacKeyStr;
    private static final String MAC_ALGORITHM = "HmacMD5";


    public static void main(String[] args) throws IOException {
        ClientThread clientThread1 = new ClientThread();
        ClientThread clientThread2 = new ClientThread();
        clientThread1.start();
        clientThread2.start();
    }

    public void start() throws IOException {
        this.publicKeyMap = new HashMap<>();
        setInitComponents();
    }


    @Override
    public void run() {
        try {
            this.start();
        } catch (IOException e) {
            e.printStackTrace();
        }
    }


    public void setInitComponents() {
        jFrame = new JFrame("安全聊天系统");
        jFrame.setFont(MY_FONT);
        jFrame.setSize(350, 200);
        jFrame.setDefaultCloseOperation(JFrame.EXIT_ON_CLOSE);
        //设置JPanel
        JPanel panel = new JPanel();
        panel.setBounds(0, 0, 350, 200);
        jFrame.add(panel);
        placeInitComponents(panel);
        jFrame.setVisible(true);
    }

    private void placeInitComponents(JPanel panel) {

        //端口标签
        JLabel portLabel = new JLabel("指定本机端口:");
        portLabel.setBounds(10, 20, 80, 25);
        panel.add(portLabel);

        // 输入端口的文本域
        final JTextField portText = new JTextField(20);
        portText.setBounds(80, 20, 125, 25);
        panel.add(portText);


        JButton clientButton = new JButton("启动本机聊天");
        clientButton.setBackground(Color.ORANGE);
        clientButton.setBounds(135, 80, 100, 25);
        clientButton.setVerticalAlignment(SwingConstants.CENTER);
        panel.add(clientButton);

        //跳转至客户端界面的ActionListener事件
        clientButton.addActionListener(new ActionListener() {
            @Override
            public void actionPerformed(ActionEvent e) {
                try {
                    myPort = Integer.parseInt(portText.getText());
                    serverSocket = new ServerSocket(myPort);
                    System.out.println("客户端已启动聊天服务器，端口为: " + myPort + "\n");
                    if (login()) {
                        JOptionPane.showMessageDialog(panel, "服务器启动成功", "SUCCESS", JOptionPane.INFORMATION_MESSAGE);
                    } else {
                        JOptionPane.showMessageDialog(panel, "系统错误！请重新登录", "ERROR", JOptionPane.WARNING_MESSAGE);
                    }
                } catch (IOException ex) {
                    ex.printStackTrace();
                }
            }
        });
        JButton keyButton = new JButton("生成本机密钥对");
        keyButton.setBackground(Color.PINK);
        keyButton.setBounds(135, 150, 100, 25);
        keyButton.setVerticalAlignment(SwingConstants.BOTTOM);
        panel.add(keyButton);

        //跳转至客户端界面的ActionListener事件
        keyButton.addActionListener(new ActionListener() {
            @Override
            public void actionPerformed(ActionEvent e) {
                if (serverSocket == null) {
                    JOptionPane.showMessageDialog(panel, "请先开启服务器！", "ERROR", JOptionPane.WARNING_MESSAGE);
                } else {
                    try {
                        myKeyPair = RSAUtils.generateKeyPair();
//                        test();
                        System.out.println("生成完毕，开始传输" + "\n");
                        boolean result = updateKey();
                        if (result) {
                            JOptionPane.showMessageDialog(panel, "密钥对生成成功！\n已上传至服务器", "SUCCESS", JOptionPane.INFORMATION_MESSAGE);
                        } else {
                            JOptionPane.showMessageDialog(panel, "系统错误！", "ERROR", JOptionPane.WARNING_MESSAGE);
                        }
                    } catch (Exception ex) {
                        ex.printStackTrace();
                    }
                }
            }
        });

        JButton talkButton = new JButton("进入聊天器");
        talkButton.setBackground(Color.GREEN);
        talkButton.setBounds(135, 180, 100, 25);
        talkButton.setVerticalAlignment(SwingConstants.BOTTOM);
        panel.add(talkButton);

        //跳转至客户端界面的ActionListener事件
        talkButton.addActionListener(new ActionListener() {
            @Override
            public void actionPerformed(ActionEvent e) {
                if (serverSocket != null && myKeyPair != null) {
                    jFrame.dispose();
                    initFriendFrame();
                } else {
                    JOptionPane.showMessageDialog(panel, "请启动聊天服务器并完成密钥对的生成！", "ERROR", JOptionPane.WARNING_MESSAGE);
                }
            }
        });
    }

    private void initFriendFrame() {
        jFrame = new JFrame("安全聊天系统" + " 本机: " + myPort);
        jFrame.setFont(MY_FONT);
        jFrame.setSize(350, 200);
        jFrame.setDefaultCloseOperation(JFrame.EXIT_ON_CLOSE);

        JPanel panel = new JPanel();
        panel.setBounds(0, 0, 350, 200);
        jFrame.add(panel);
        placeFriendComponents(panel);
        jFrame.setVisible(true);
    }

    private void placeFriendComponents(JPanel panel) {
        //端口标签
        JLabel portLabel = new JLabel("输入好友的端口号:");
        portLabel.setBounds(10, 25, 80, 25);
        panel.add(portLabel);

        // 输入端口的文本域
        final JTextField portText = new JTextField(20);
        portText.setBounds(80, 20, 125, 25);
        panel.add(portText);

        JButton friendButton = new JButton("进入聊天框");
        friendButton.setBackground(Color.ORANGE);
        friendButton.setBounds(135, 80, 100, 25);
        friendButton.setVerticalAlignment(SwingConstants.CENTER);
        panel.add(friendButton);

        //跳转至客户端界面的ActionListener事件
        friendButton.addActionListener(e -> {
            try {
                friendPort = Integer.valueOf(portText.getText());
                System.out.println("好友ID: " + friendPort + "\n");
                boolean result = checkOnline(friendPort);
                if (result) {
                    JOptionPane.showMessageDialog(panel, "对方在线！\n正在交换密钥", "SUCCESS", JOptionPane.INFORMATION_MESSAGE);
                    getKey(friendPort);
                    System.out.println("获取到了对方的公钥，存入自身的公钥库" + "\n" + publicKeyMap.get(friendPort) + "\n");

                    if (!getMacKey(friendPort)) {
                        myMacKeyStr = "266f5fe18e714688a083df4ca9f78064";//MACUtils.generateKey()
                        System.out.println("对方未生成MAC密钥，由本机 " + myPort + " 生成了会话MAC密钥: " + myMacKeyStr + "\n");
                        uploadMacKey(myMacKeyStr);
                    }


                    jFrame.dispose();
                    initTalkFrame();
                } else {
                    JOptionPane.showMessageDialog(panel, "系统错误！", "ERROR", JOptionPane.WARNING_MESSAGE);
                }
            } catch (Exception ex) {
                ex.printStackTrace();
            }
        });
    }

    private void initTalkFrame() {
        jFrame = new JFrame("安全聊天系统" + " 本机: " + myPort);
        jFrame.setFont(MY_FONT);
        jFrame.setSize(350, 200);
        jFrame.setDefaultCloseOperation(JFrame.EXIT_ON_CLOSE);

        JPanel panel = new JPanel();
        panel.setBounds(0, 0, 350, 200);
        jFrame.add(panel);
        placeTalkComponents(panel);
        jFrame.setVisible(true);
    }

    private void placeTalkComponents(JPanel panel) {
        JTextArea jTextArea = new JTextArea(combineContent());
        jTextArea.setColumns(5);
        jTextArea.setLineWrap(true);
        jTextArea.setBackground(Color.lightGray);
        jTextArea.setBounds(10, 20, 300, 300);
        jTextArea.setFont(new Font("楷体", Font.BOLD, 16));
        panel.add(jTextArea);


        Thread listenThread = new Thread(() -> {
            while (true) {
                Socket socket = null;
                InputStream inputStream = null;
                OutputStream outputStream = null;
                try {
                    socket = serverSocket.accept();
                } catch (IOException e) {
                    e.printStackTrace();
                }
                try {
                    inputStream = socket.getInputStream();
                    outputStream = socket.getOutputStream();
                } catch (IOException e) {
                    e.printStackTrace();
                }

                byte[] bytes = new byte[256];
                int len = 0;
                try {
                    len = inputStream.read(bytes);
                } catch (IOException e) {
                    e.printStackTrace();
                }

                byte[] plainData = new byte[256];
                try {
                    plainData = RSAUtils.decrypt(bytes, myKeyPair.getPrivate());
                } catch (Exception e) {
                    System.out.println("消息传递解密出错" + myPort + myKeyPair.getPublic() + "\n");
                    e.printStackTrace();
                }

                String str = new String(plainData, 0, plainData.length);
                try {
                    send(outputStream, "success");
                } catch (IOException e) {
                    e.printStackTrace();
                }
                String text = str.substring(0, str.length());
                System.out.println("明文:" + text + "\n");


                byte[] macBytes = new byte[1024];
                int len2 = 0;
                try {
                    len2 = inputStream.read(macBytes);
                } catch (IOException e) {
                    e.printStackTrace();
                }
                String originalMac = new String(macBytes, 0, len2);

                //MAC验证消息完整性
//                inputStream.read()
                Key key = new SecretKeySpec(myMacKeyStr.getBytes(), "");
                try {
                    String nowMac = new String(MACUtils.mac("HmacMD5", key, text.getBytes(StandardCharsets.UTF_8)));
                    boolean tag = originalMac.equals(nowMac);
                    if (tag) {
                        send(outputStream, "success");
                        System.out.println("最新一条消息的消息完整性（MAC）无误！\n");
                    } else {
                        send(outputStream, "false");
                    }
                    System.out.println("收到消息的MAC值:" + nowMac);
                } catch (Exception e) {
                    e.printStackTrace();
                }
                updateContent(jTextArea, text, false);
            }
        });
        listenThread.start();


        //ip地址标签
        JLabel userLabel = new JLabel("文本内容:");
        userLabel.setBounds(10, 20, 80, 25);
        panel.add(userLabel);

        //输入ip地址的文本域
        JTextField messageText = new JTextField(20);
        messageText.setBounds(100, 20, 165, 25);
        panel.add(messageText);


        JButton sendButton = new JButton("发送");
        sendButton.setBackground(Color.GRAY);
        sendButton.setBounds(10, 20, 100, 25);
        sendButton.setVerticalAlignment(SwingConstants.BOTTOM);
        panel.add(sendButton);


        //发送消息按键的ActionListener事件
        sendButton.addActionListener(e -> {
            if (talkWithFriends(messageText.getText())) {
                updateContent(jTextArea, messageText.getText(), true);
            } else {
                JOptionPane.showMessageDialog(panel, "系统错误！", "ERROR", JOptionPane.WARNING_MESSAGE);
            }

        });


        JButton exitButton = new JButton("退出");
        exitButton.setBackground(Color.GRAY);
        exitButton.setBounds(10, 20, 100, 25);
        exitButton.setVerticalAlignment(SwingConstants.BOTTOM);
        panel.add(exitButton);

        //跳转至客户端界面的ActionListener事件
        exitButton.addActionListener(new ActionListener() {
            @Override
            public void actionPerformed(ActionEvent e) {
                try {
                    logout();
                    JOptionPane.showMessageDialog(panel, "您的账号已离线", "SUCCESS", JOptionPane.INFORMATION_MESSAGE);
                } catch (IOException ex) {
                    ex.printStackTrace();
                }
                System.exit(0);
            }
        });


    }

    private void updateContent(JTextArea jTextArea, String str, boolean isMe) {
        if (isMe) {
            talkContent[contentTag % 5] = friendPort + ": " + str;
        } else {
            talkContent[contentTag % 5] = myPort + ": " + str;
        }

        contentTag++;
        jTextArea.setText(combineContent());
    }


    public boolean talkWithFriends(String content) {
        Socket socket;
        try {
            socket = new Socket(TARGET_IP, friendPort);
            InputStream inputStream = socket.getInputStream();
            OutputStream outputStream = socket.getOutputStream();
            byte[] bytes = content.getBytes();

            byte[] cipher = RSAUtils.encrypt(bytes, publicKeyMap.get(friendPort));
            //适用RSAUtils加密文本内容
            outputStream.write(cipher, 0, cipher.length);

            System.out.println("使用对方公钥加密: " + friendPort + "\n加密密文:" + new String(cipher) + "\n");

            System.out.println("加密密钥: " + publicKeyMap.get(friendPort) + "\n");

            Key key = new SecretKeySpec(myMacKeyStr.getBytes(), "");
            String mac = new String(MACUtils.mac("HmacMD5", key, content.getBytes(StandardCharsets.UTF_8)));
            System.out.println("发送消息的mac值:  " + mac);
            String answer = receive(inputStream);
            outputStream.write(mac.getBytes(StandardCharsets.UTF_8));
            answer = receive(inputStream);
            return answer.equals("success");
        } catch (Exception e) {
            e.printStackTrace();
            return false;
        }
    }


    /**
     * 简化的接收一条简短文本的方法
     */
    public String receive(InputStream inputStream) {
        String str = null;
        try {
            byte[] bytes = new byte[256];
            int len = inputStream.read(bytes);
            str = new String(bytes, 0, len).trim();
        } catch (Exception e) {
            e.printStackTrace();
        }
        return str;
    }

    public void send(OutputStream outputStream, String message) throws IOException {
        outputStream.write(message.getBytes(StandardCharsets.UTF_8));
    }


    /**
     * 初始化以及更新自己账户生成的密钥对中的公钥
     * 将其发布至公钥服务器供好友取用
     */
    public boolean updateKey() {

        Socket socket;
        try {
            socket = new Socket(SERVER_IP, SERVER_PORT);
            OutputStream outputStream = socket.getOutputStream();
            InputStream inputStream = socket.getInputStream();

            //发送头部消息指明操作
            send(outputStream, "updateKey");
            System.out.print("头部消息发送完毕");
            System.out.println(receive(inputStream) + "\n");

            send(outputStream, String.valueOf(myPort));


            //重新建立Socket，以发送序列化公钥数据
            socket = new Socket(SERVER_IP, SERVER_PORT);
            inputStream = socket.getInputStream();
            outputStream = socket.getOutputStream();

            ObjectOutputStream objectOutputStream = new ObjectOutputStream(outputStream);
            PublicKey publicKey = myKeyPair.getPublic();
            objectOutputStream.writeObject(publicKey);
            System.out.println("对象发完了");

            System.out.println(myKeyPair.getPublic() + "\n");
//            socket.shutdownOutput();
            String answer = receive(inputStream);
            System.out.println();
            return answer.equals("success");
        } catch (IOException e) {
            e.printStackTrace();
            System.out.println("更新密钥失败！！！！");
            return false;
        }
    }


    /**
     * 获取公钥服务器中的目标公钥用于非对称加密数据进行交流
     */
    public boolean getKey(Integer targetPort) {
        Socket socket;
        try {
            socket = new Socket(SERVER_IP, SERVER_PORT);
            OutputStream outputStream = socket.getOutputStream();
            InputStream inputStream = socket.getInputStream();

            //发送头部消息指明操作
            send(outputStream, "getKey");
            System.out.print("头部消息发送完毕");
            System.out.println(" :" + receive(inputStream) + "\n");

            //发送需要的端口
            send(outputStream, String.valueOf(targetPort));


            //重新建立Socket，以接收序列化公钥数据
            socket = new Socket(SERVER_IP, SERVER_PORT);
            inputStream = socket.getInputStream();
            ObjectInputStream objectInputStream = new ObjectInputStream(inputStream);

            //接收公钥
            PublicKey publicKey = (PublicKey) objectInputStream.readObject();
            System.out.println("获取到了 " + targetPort + " 的公钥" + "\n");
            publicKeyMap.put(targetPort, publicKey);
            return true;
        } catch (Exception e) {
            e.printStackTrace();
            return false;
        }
    }


    /**
     * 检查好友是否在线
     *
     * @param port 好友ID
     * @return 成功与否
     */
    public boolean checkOnline(Integer port) throws IOException {
        Socket socket = new Socket(SERVER_IP, SERVER_PORT);
        OutputStream outputStream = socket.getOutputStream();
        InputStream inputStream = socket.getInputStream();
        send(outputStream, "checkOnline");
        System.out.print("头部消息发送完毕");
        System.out.println(" :" + receive(inputStream) + "\n");

        send(outputStream, String.valueOf(friendPort));
        String answer = receive(inputStream);
        return answer.equals("success");
    }


    /**
     * 用于登录此账户，将其状态同步至服务器存储
     */
    public boolean login() throws IOException {
        Socket socket = new Socket(SERVER_IP, SERVER_PORT);
        OutputStream outputStream = socket.getOutputStream();
        InputStream inputStream = socket.getInputStream();
        send(outputStream, "login");
        System.out.print("头部消息发送完毕");
        System.out.println(" :" + receive(inputStream) + "\n");

        send(outputStream, String.valueOf(myPort));
        String answer = receive(inputStream);
        return answer.equals("success");
    }

    /**
     * 用于登出账户
     */
    public void logout() throws IOException {
        Socket socket = new Socket(SERVER_IP, SERVER_PORT);
        OutputStream outputStream = socket.getOutputStream();
        InputStream inputStream = socket.getInputStream();
        send(outputStream, "logout");
        System.out.print("头部消息发送完毕");
        System.out.println(" :" + receive(inputStream) + "\n");

        send(outputStream, String.valueOf(myPort));
        String answer = receive(inputStream);
    }


    public boolean getMacKey(Integer port) throws IOException {
        Socket socket = new Socket(SERVER_IP, SERVER_PORT);
        OutputStream outputStream = socket.getOutputStream();
        InputStream inputStream = socket.getInputStream();
        send(outputStream, "getMacKey");
        System.out.print("头部消息发送完毕");
        System.out.println(" :" + receive(inputStream) + "\n");
        send(outputStream, String.valueOf(port));

        if (!receive(inputStream).equals("success")) {
            return false;
        }
        socket = new Socket(SERVER_IP, SERVER_PORT);
        inputStream = socket.getInputStream();

        byte[] bytes = new byte[256];
        int len = inputStream.read(bytes);

//        byte[] plainData = new byte[256];

        //数据签名验证签名
        String answer = "";
        try {
            answer = RC4Utils.HLoveyRC4(new String(bytes).trim(), RC4_KEY);
            System.out.println(answer);
//            plainData = RSAUtils.decrypt(bytes, publicKeyMap.get(friendPort));
        } catch (Exception e) {
            System.out.println("MAC解密出现错误 " + myPort + "\n");
            e.printStackTrace();
        }
        myMacKeyStr = answer;
        System.out.println("收到密钥:" + myMacKeyStr + "\n");
        return true;
    }


    private boolean uploadMacKey(String key) throws Exception {
        Socket socket = new Socket(SERVER_IP, SERVER_PORT);
        OutputStream outputStream = socket.getOutputStream();
        InputStream inputStream = socket.getInputStream();
        send(outputStream, "uploadMacKey");
        System.out.print("头部消息发送完毕");
        System.out.println(" :" + receive(inputStream));
        outputStream.write(String.valueOf(myPort).getBytes());
        System.out.println("端口上传:  " + receive(inputStream) + "\n");
        //加密MAC密钥
//        byte[] cipher = RSAUtils.encrypt(key.getBytes(), myKeyPair.getPublic());
//        String str = new String(cipher, 0, cipher.length);

//        System.out.println("加密MAC密钥成功 "+new String(cipher));
        String temp = RC4Utils.HLoveyRC4(key,RC4_KEY);
        send(outputStream, temp);
        if (receive(inputStream).equals("success")) {
            System.out.println("用户 " + myPort + " mac密钥上传成功" + "\n");
            return true;
        } else {
            return false;
        }
    }

    private String combineContent() {
        StringBuilder sb = new StringBuilder();
        for (String s : talkContent) {
            sb.append(s).append("\n");
        }
        return sb.toString();
    }
}
